package com.zyx.javademo.collection.set;

import com.zyx.javademo.bean.User;
import com.zyx.javademo.collection.queue.PriorityQueueTopK;

import java.math.BigDecimal;
import java.util.*;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/**
 * @author Yaxi.Zhang
 * @since 2021/5/28 12:05
 * desc:
 */
public class TreeSetTopK<T> {
    /**
     * TreeSet的边界，Top K 问题中的 K
     */
    private final int boundary;

    /**
     * 用来构造一个有界的
     */
    private final TreeSet<T> boundaryTreeSet;

    /**
     * 通过自定义边界 boundary 可以求解 top K 问题
     * 通过自定义比较器 comparator 可以控制求解 top K 大 还是 top K 小
     * @param boundary 边界 K
     * @param comparator 数据比较器
     */
    public TreeSetTopK(int boundary, Comparator<T> comparator) {
        this.boundary = boundary;
        boundaryTreeSet = new TreeSet<>(comparator);
    }

    /**
     * 求解数据流中的top K， 将结果写入List中
     * @param dataStream 数据流
     * @param results top K 结果
     */
    public void topK(Stream<T> dataStream, List<T> results) {
        dataStream.forEach(this::add);
        results.addAll(boundaryTreeSet);
    }


    /**
     * 向有界堆中添加元素的帮助方法
     * @param t 待添加数据
     */
    private void add(T t) {
        boundaryTreeSet.add(t);
        if (boundaryTreeSet.size() > boundary) {
            boundaryTreeSet.pollLast();
        }
    }

    public static void main(String[] args) {
        // perfDemo();
        // multiDemo();
        udComDemo();
    }


    public static void udComDemo() {
        User user1 = new User("zhangsan", 18, new BigDecimal(180));
        User user2 = new User("lisi", 22, new BigDecimal(180));
        User user3 = new User("wangwu", 20, new BigDecimal(180));
        User user4 = new User("zhaoliu", 21, new BigDecimal(180));
        User user5 = new User("tianqi", 23, new BigDecimal(180));
        User user6 = new User("wangba", 21, new BigDecimal(180));
        TreeSetTopK<User> topK = new TreeSetTopK<User>(4, (o1, o2) -> {
            int res = Integer.compare(o2.getAge(), o1.getAge());
            return res == 0 ? 1 : res;
        });
        List<User> list = new ArrayList<>();
        list.add(user1);
        list.add(user2);
        list.add(user3);
        list.add(user4);
        list.add(user5);
        list.add(user6);
        Stream<User> stream = StreamSupport.stream(Spliterators.spliteratorUnknownSize(
                list.iterator(), Spliterator.ORDERED), false);
        List<User> resList = new ArrayList<>();
        topK.topK(stream, resList);
        System.out.println(resList);
    }

    /**
     * 测试TreeSet求TopK的性能
     */
    public static void perfDemo() {
        // 构造一个 范围为 [0, 2^30] 的 Integer 流，通过limit可以控制大小
        final int upLimit = 1 << 30;
        Stream<Integer> stream = Stream.generate(Math::random)
                .map(d -> d * upLimit)
                .map(d -> (int) Math.round(d))
                .limit(100_000_000);

        // 将 (o1, o2) -> (o1 - o2) 换成 (o1, o2) -> (o2 - o1) 可以求解 top K 小
        TreeSetTopK<Integer> topK = new TreeSetTopK<>(10, (o1, o2) -> (o2 - o1));

        List<Integer> results = new ArrayList<>();
        long startTime = System.currentTimeMillis();
        topK.topK(stream, results);

        long endTime = System.currentTimeMillis();
        System.out.println("results: " + results);
        System.out.println("cost: " + (endTime - startTime) / 1000.0);
    }

    public static void multiDemo() {
        Iterable<Integer> iterable = Arrays.asList(1, 2, 2, 3, 6, 8, 7, 9, 10, 9);
        // Iterable -> Spliterators -> Stream -> List
        Stream<Integer> stream = StreamSupport.stream(iterable.spliterator(), false);
        TreeSetTopK<Integer> topK = new TreeSetTopK<>(5, (o1, o2) -> (o1 - o2));
        List<Integer> results = new ArrayList<>();
        topK.topK(stream, results);
        System.out.println("results: " + results);
    }
}
